#!/usr/bin/env python3

import os, sys
from astropy.io import fits
import numpy as np

def remove_overscan(imgdata):
    img = np.zeros((9232,9216), dtype=np.int16)
#    for i in range(4):
#        img[0:4616, i*1152:(i+1)*1152] = imgdata[0:4616,(i*1250+27):((i+1)*1250-71)] - imgdata[0:4616,((i+1)*1250-71):((i+1)*1250)].mean() 
#        img[4616:, i*1152:(i+1)*1152] = imgdata[4784:,(i*1250+27):((i+1)*1250-71)] - imgdata[4784:,((i+1)*1250-71):((i+1)*1250)].mean() 
#    for i in range(4,8):
#        img[0:4616, i*1152:(i+1)*1152] = imgdata[0:4616,(i*1250+71):((i+1)*1250-27)] - imgdata[0:4616,(i*1250):(i*1250+71)].mean() 
#        img[4616:, i*1152:(i+1)*1152] = imgdata[4784:,(i*1250+71):((i+1)*1250-27)] - imgdata[4784:,(i*1250):(i*1250+71)].mean()

    for i in range(4):
        img[0:4616, i*1152:(i+1)*1152] = imgdata[0:4616,(i*1250+27):((i+1)*1250-71)]
        img[4616:, i*1152:(i+1)*1152] = imgdata[4784:,(i*1250+27):((i+1)*1250-71)] 
        oc1 = np.median(imgdata[0:4616,((i+1)*1250-71):((i+1)*1250)], axis=1).astype(float) 
        oc2 = np.median(imgdata[4784:,((i+1)*1250-71):((i+1)*1250)], axis=1).astype(float)
        print('oc1.shape: {}'.format(oc1.shape))
        for j in range(4616):
            img[j,i*1152:(i+1)*1152] = img[j,i*1152:(i+1)*1152] - oc1[j]
            img[4616+j,i*1152:(i+1)*1152] = img[4616+j,i*1152:(i+1)*1152]- oc2[j]
    for i in range(4,8):
        img[0:4616, i*1152:(i+1)*1152] = imgdata[0:4616,(i*1250+71):((i+1)*1250-27)]
        img[4616:, i*1152:(i+1)*1152] = imgdata[4784:,(i*1250+71):((i+1)*1250-27)]
        oc1 = np.median(imgdata[0:4616,(i*1250):(i*1250+71)], axis=1).astype(float)
        oc2 = np.median(imgdata[4784:,(i*1250):(i*1250+71)], axis=1).astype(float)
        for j in range(4616):
            img[j, i*1152:(i+1)*1152] = img[j, i*1152:(i+1)*1152] - oc1[j]
            img[4616+j,i*1152:(i+1)*1152] = img[4616+j,i*1152:(i+1)*1152] - oc2[j]

    return img


def run(imglist):
    pass



def main():
    fname = sys.argv[1]
    imgdata = fits.getdata(fname).astype(float)
    img = remove_overscan(imgdata)
    
    hdu = fits.PrimaryHDU(img.astype(np.float32))
    fout = fname.replace('.fits', '_overscan_removed.fits')
    hdu.writeto(fout, overwrite=True)

    cmd = 'ds9 {} &'.format(fout)
    print('running system cmd: {}'.format(cmd))
    os.system(cmd)

if __name__ == '__main__':
    main()
